#include "StdAfx.h"
#include "NetworkUtils.h"

#include <vector>
#include <regex>

tchar_string NetworkUtils::getDriveConnection(tchar_string const& drive)
{
    DWORD buffSize = 0;
    if (ERROR_MORE_DATA == WNetGetConnection(drive.c_str(), 0, &buffSize))
    {
        std::vector<TCHAR> buffer(buffSize + 1);
        if (NO_ERROR == WNetGetConnection(drive.c_str(), &buffer[0], &buffSize))
        {
            return &buffer[0];
        }
    }

    return tchar_string();
}

tchar_string NetworkUtils::getComputerName()
{
    TCHAR buf[MAX_COMPUTERNAME_LENGTH + 1] = {0};
    DWORD size = ARRAYSIZE(buf);
    GetComputerName(buf, &size);
    return &buf[0];
}


NetworkUtils::SharesList NetworkUtils::getShares()
{
    SharesList result;

    CRegKey shares;
    if (ERROR_SUCCESS == shares.Open(HKEY_LOCAL_MACHINE, _T("SYSTEM\\CurrentControlSet\\Services\\LanmanServer\\Shares"), KEY_READ))
    {
        std::tr1::basic_regex<TCHAR> shareName(_T("^ShareName=(.*)$"));
        std::tr1::basic_regex<TCHAR> sharePath(_T("^Path=(.*)$"));

        std::vector<TCHAR> valueName(16*1024);
        DWORD valueNameLength = valueName.size();
        DWORD rgType = 0;
        std::vector<TCHAR> data(16*1024);
        DWORD dataSize = data.size();
        for (long i = 0; ERROR_NO_MORE_ITEMS != RegEnumValue(shares.m_hKey, i, &valueName[0], &valueNameLength, 0, &rgType, (LPBYTE)&data[0], &dataSize); ++i)
        {
            std::replace(data.begin(), data.begin() + dataSize, TCHAR(_T('\0')), TCHAR(_T('\n')));
            tchar_string shareSpec((TCHAR*)&data[0]);
            std::tr1::match_results<tchar_string::const_iterator> shareNameRes;
            std::tr1::match_results<tchar_string::const_iterator> sharePathRes;
            if (std::tr1::regex_search(shareSpec, shareNameRes, shareName) && 
                std::tr1::regex_search(shareSpec, sharePathRes, sharePath))
            {
                tchar_string name(shareNameRes[1]);
                tchar_string path(sharePathRes[1]);
                result[name] = path;
            }

            dataSize = data.size();
            valueNameLength = valueName.size();
        }
    }

    return result;
}